import copy
import numpy as np
# 明确的知道每个状态的回报，即R和S是已知的，状态转移矩阵P也是已知的
class CliffWalkingEnv:
    """ 悬崖漫步环境"""
    def __init__(self, ncol=12, nrow=4):
        self.ncol = ncol  # 定义网格世界的列
        self.nrow = nrow  # 定义网格世界的行
        # 转移矩阵P[state][action] = [(p, next_state, reward, done)]包含下一个状态和奖励
        self.P = self.createP()

    def createP(self):
        """
        返回大小为48的列表, 用索引来表示状态, 每个索引下有4个值, 表示4个动作，上下左右，用元组来表示
        个人觉得只需要二维数组即可，并不需要下面的三维，例如：[ [(),(),(),()], [(),(),(),()] ]
        但是为什么这里是三维的呢？因为这里的一个动作只对应一个状态，当然了，如果不是悬崖的话，一个动作就可能对应多个状态，所以还是三维合理
        [ [[()],[()],[()],[()]], [[()],[()],[()],[()]],    ]
        假设存在的可能为：[ [[(),()],[()],[(),(),()],[()]], [[(),()],[()],[(),(),()],[()] ]
        元组里面为4个数组，分别表示(p, next_state, reword, done)，其中p表示走向下一个状态的概率，
        因为是悬崖，一次走一个，所以确定了方向之后，下一个状态一定是确定了，比如从索引0向下走，next_state一定是索引12
        """
        # 初始化
        P = [[[] for j in range(4)] for i in range(self.nrow * self.ncol)]
        # 4种动作, change[0]:上,change[1]:下, change[2]:左, change[3]:右。坐标系原点(0,0)
        # 为什么这么定义呢？因为对于上下移动来说，x的变化为0，上的话就是y=y-1，下的话就是y=y+1，我们以点的变化（change）作为值
        # 上就是[0, -1]，下就是[0, 1]，同理左右移动时，y不变（变化为0），左的话x=x-1，右的话x=x+1，所以
        # 左就是[-1, 0]，右就是[1, 0]
        # 定义在左上角，我们默认以矩阵的左上角所谓起点，也就是(0, 0)
        change = [[0, -1], [0, 1], [-1, 0], [1, 0]]
        for i in range(self.nrow):
            for j in range(self.ncol):
                # 对一个索引（状态）的四个方向（动作）设置值，并append到状态转移矩阵列表P中
                for a in range(4):
                    # 位置（出发点）在悬崖或者目标状态,因为无法继续交互,任何动作奖励都为0
                    # 我们设定的是矩阵的最后一行的中间全部是悬崖，最后一行的最后一列是终点
                    if i == self.nrow - 1 and j > 0:
                        P[i * self.ncol + j][a] = [(1, i * self.ncol + j, 0,
                                                    True)]
                        continue
                    # 依据上面关于change的上下左右定义，next_x = j + change[a][0]的，但是考虑到当前状态可能为最左（0），或最右(ncol-1)，只有在中间才需要变动
                    next_x = min(self.ncol - 1, max(0, j + change[a][0]))
                    next_y = min(self.nrow - 1, max(0, i + change[a][1]))
                    next_state = next_y * self.ncol + next_x
                    reward = -1
                    done = False
                    # 下一个位置在悬崖或者终点（悬崖和终点肯定是最后一行，且不是第一格）
                    if next_y == self.nrow - 1 and next_x > 0:
                        done = True
                        if next_x != self.ncol - 1:  # 下一个位置在悬崖
                            reward = -100
                    P[i * self.ncol + j][a] = [(1, next_state, reward, done)]
        return P

# 我们要注意区分变量P和变量pi，一个是状态转移矩阵，一个是某个状态下的动作选择概率，P是概率probability，而pi是π的含义
class PolicyIteration:
    """ 策略迭代算法 """
    def __init__(self, env: CliffWalkingEnv, theta: float, gamma: float):
        self.env = env
        self.v = [0] * self.env.ncol * self.env.nrow  # 初始化价值为0，初始值为[0,0,...0]，大小为48
        # 初始化为均匀随机策略，初始值为 [ [0.25, 0.25, 0.25, 0.25]，[0.25, 0.25, 0.25, 0.25]，...[0.25, 0.25, 0.25, 0.25] ]
        # 大小为48，其中里面的[0.25, 0.25, 0.25, 0.25]，表示当前状态（索引）执行上下左右4个动作的概率，默认都是1/4
        self.pi = [[0.25, 0.25, 0.25, 0.25] for i in range(self.env.ncol * self.env.nrow)]  
        self.theta = theta  # 策略评估收敛阈值
        self.gamma = gamma  # 折扣因子
    # 策略评估的过程，其本质就是不断计算每个状态的价值的过程，且更新价值的过程（策略不变），状态价值=下面的动作价值之和
    def policy_evaluation(self):  
        cnt = 1  # 计数器
        while 1:
            max_diff = 0
            # 迭代后的价值状态函数（列表，V(S)）
            new_v: list[float] = [0] * self.env.ncol * self.env.nrow
            for s in range(self.env.ncol * self.env.nrow):
                qsa_list: float = []  # 开始计算某状态s下的所有Q(s,a)价值，大小为4
                for a in range(4):
                    # 对每个状态下的4个动作，取其reward，并加权平均
                    qsa = 0
                    # 其实这里的res大小是1，for循环只会有一次，因为一个动作只会对应一个状态（概率=1）
                    for res in self.env.P[s][a]:
                        p, next_state, r, done = res
                        # 本章环境比较特殊,奖励和下一个状态有关,所以需要和状态转移概率相乘
                        qsa += p * (r + self.gamma * self.v[next_state] * (1 - done))
                    qsa_list.append(self.pi[s][a] * qsa)
                # 对某状态下的4个Q求和，就是迭代后的状态价值
                new_v[s] = sum(qsa_list)  # 状态价值函数和动作价值函数之间的关系
                # 从这里可以看出max_diff在for循环之内，表示新状态列表和就状态列表之间每个对应索引状态的差值，取最大值
                # 如果连这个最大差值都达标了，那就说明剩下的差值肯定达标了，新状态就是想要的，停止迭代，否则继续迭代
                max_diff = max(max_diff, abs(new_v[s] - self.v[s]))
            # 更新当前状态，因为每次while进来之后，new_v都会用上self.v，就是说是基于上一次的变更，具体是哪些变更了呢？就是V(S)，状态价值
            # 而我们的目的就是更新状态价值函数，其他的不变，比如动作，及其动作概率，是不可能变的
            self.v = new_v
            if max_diff < self.theta: 
                break  # 满足收敛条件,退出评估迭代
            cnt += 1
        print("策略评估进行%d轮后完成" % cnt)

    # 策略提升的过程，也就是更新pi（π）的一个过程
    def policy_improvement(self):  
        # 遍历每个状态，也就是每一个格子
        for s in range(self.env.nrow * self.env.ncol):
            # 计算每个状态的Q(s, a)
            qsa_list = []
            for a in range(4):
                qsa = 0
                for res in self.env.P[s][a]:
                    p, next_state, r, done = res
                    qsa += p * (r + self.gamma * self.v[next_state] * (1 - done))
                qsa_list.append(qsa)
            # 取4个动作中，动作价值Q最大的一个
            maxq = max(qsa_list)
            # 计算有几个动作得到了最大的Q值，因为这里的策略提升，就是重新分配概率
            # 我们只留下Q最大的几个，并平分概率，其他小于Q的，概率设为0，但是保留，并不会删掉，不然就影响上下左右的布局了
            cntq = qsa_list.count(maxq)  
            # 让这些动作均分概率，这里又是Python语法的骚操作了，不喜欢的可以用for循环写，更加直观
            self.pi[s] = [1 / cntq if q == maxq else 0 for q in qsa_list]
        print("策略提升完成")
        return self.pi

    def policy_iteration(self):  # 策略迭代
        while 1:
            self.policy_evaluation()
            # 将列表进行深拷贝,方便接下来进行比较，浅拷贝和深拷贝是有区别的
            old_pi = copy.deepcopy(self.pi) 
            new_pi = self.policy_improvement()
            # 本人很好奇，什么时候等式成立，是不是每个状态下的4个概率，要么最大，要么=0，比如 [[0,0.333,0.333,0.333],[0,0,0,1],[0.5,0.5,0,0]]
            if old_pi == new_pi: 
                break

def print_agent(agent: PolicyIteration, action_meaning, disaster=[], end=[]):
    print("状态价值：")
    for i in range(agent.env.nrow):
        for j in range(agent.env.ncol):
            # 为了输出美观,保持输出6个字符，这个写法太不友好，我觉得应该改，不就是输出宽度为6，小数点占3位嘛
            # print('%6.6s' % ('%.3f' % agent.v[i * agent.env.ncol + j]), end=' ')
            value = agent.v[i * agent.env.ncol + j]
            print(f'{value:6.3f}', end=' ')
        print()

    print("策略：")
    for i in range(agent.env.nrow):
        for j in range(agent.env.ncol):
            # 一些特殊的状态,例如悬崖漫步中的悬崖
            if (i * agent.env.ncol + j) in disaster:
                print('****', end=' ')
            elif (i * agent.env.ncol + j) in end:  # 目标状态
                print('EEEE', end=' ')
            else:
                # 这里的a的值类似于(0, 0, 0.5, 0.5)，表示四个方向的概率
                a = agent.pi[i * agent.env.ncol + j]
                pi_str = ''
                for k in range(len(action_meaning)):
                    pi_str += action_meaning[k] if a[k] > 0 else 'o'
                print(pi_str, end=' ')
        print()


env = CliffWalkingEnv()
action_meaning = ['^', 'v', '<', '>']
theta = 0.001
gamma = 0.9
agent = PolicyIteration(env, theta, gamma)
agent.policy_iteration()
print_agent(agent, action_meaning, list(range(37, 47)), [47])

class ValueIteration:
    """ 价值迭代算法 """
    def __init__(self, env, theta, gamma):
        self.env = env
        self.v = [0] * self.env.ncol * self.env.nrow  # 初始化价值为0
        self.theta = theta  # 价值收敛阈值
        self.gamma = gamma
        # 价值迭代结束后得到的策略，为什么这里使用None，而不是[0.25,0.25,0.25,0.25]
        self.pi = [None for i in range(self.env.ncol * self.env.nrow)]

    def value_iteration(self):
        cnt = 0
        while 1:
            max_diff = 0
            new_v = [0] * self.env.ncol * self.env.nrow
            for s in range(self.env.ncol * self.env.nrow):
                qsa_list:list[float] = []  # 开始计算状态s下的所有Q(s,a)价值
                for a in range(4):
                    qsa = 0
                    for res in self.env.P[s][a]:
                        p, next_state, r, done = res
                        qsa += p * (r + self.gamma * self.v[next_state] * (1 - done))
                    qsa_list.append(qsa)  # 这一行和下一行代码是价值迭代和策略迭代的主要区别
                # 新的状态价值，直接取了最大的动作价值，抛弃了另外三个
                new_v[s] = max(qsa_list)
                max_diff = max(max_diff, abs(new_v[s] - self.v[s]))
            self.v = new_v
            # 满足收敛条件,退出评估迭代
            if max_diff < self.theta:
                break
            cnt += 1
        print("价值迭代一共进行%d轮" % cnt)
        self.get_policy()

    # 这也是更新策略pi的过程，只执行一次
    def get_policy(self):  # 根据价值函数导出一个贪婪策略
        for s in range(self.env.nrow * self.env.ncol):
            qsa_list = []
            for a in range(4):
                qsa = 0
                for res in self.env.P[s][a]:
                    p, next_state, r, done = res
                    qsa += p * (r + self.gamma * self.v[next_state] * (1 - done))
                qsa_list.append(qsa)
            maxq = max(qsa_list)
            # 计算有几个动作得到了最大的Q值，与上面的策略迭代更新原理是一样的
            cntq = qsa_list.count(maxq)
            # 让这些动作均分概率
            self.pi[s] = [1 / cntq if q == maxq else 0 for q in qsa_list]


env = CliffWalkingEnv()
action_meaning = ['^', 'v', '<', '>']
theta = 0.001
gamma = 0.9
agent = ValueIteration(env, theta, gamma)
agent.value_iteration()
print_agent(agent, action_meaning, list(range(37, 47)), [47])